.text                     # Start text section
#.balign 4                # align section to 4 bytes - which is also the default?
.global solve  	          # define global function symbol



solve:
	li t1, 0
	fcvt.s.l f0, t1			# f0 = 0.0
	mv t3, a0			# t3: total num of element
	vsetvli t0, t3, e32, m2, ta, ma # Set vector length based on 32-bit vectors
	vfmv.v.f v0, f0			# float 0.0
	mv t5, a5			# tmp base of a5
	mv t6, a6			# tmp base of a6 

loop1:
# S_vector_layer1.at(i, 0) = 0 	
# H_vector_layer1.at(i, 0) = 0
	li t2, 4 			# stride
	
	vsetvli t0, t3, e32, m2, ta, ma # Set vector length based on 32-bit vectors
	vsse32.v v0, (t5), t2		
	vsse32.v v0, (t6), t2
	
	sub t3, t3, t0			# check loop
	slli t0, t0, 2
	add t5, t5, t0
	add t6, t6, t0
	bnez t3, loop1

# ---------------------------------------------------------------------------------	
	li t1, 0			# loop: p
loop2:					# used register: t1, t4, t5, t6				
	mul s1, t1, a0
	slli s1, s1, 2
	add t4, a4, s1			# auto & x = x_testdata_window[p] // x_base + p * innode * 4
	addi t3, t1, 1			# t3 = p + 1
	mul s1, t3, a1			
	slli s1, s1, 2
	add t5, a5, s1			# S_base + (p + 1) * hidenode_1 * 4
	add t6, a6, s1			# H_base + (p + 1) * hidenode_1 * 4
					# t4(a4): x_test_data_window[p]
					# t5(a5): S_vector_layer1.line(p + 1)
					# t6(a6): H_vector_layer1.line(p + 1)
	 
	li s10, 0			# s10: j = 0
loop3:					# used register: s10, f1, f2, f3, f4
	fmv.s f1, f0			# float inGate = 0.0
	fmv.s f2, f0			# float outGate = 0.0
	fmv.s f3, f0			# float forgetGate = 0.0
	fmv.s f4, f0			# float gGate = 0.0

	li s11, 0			# m = 0, 0 <= m < innode(a0)
loop4:					# used register: s11, s2, s3, s4, s5
					# in fact, this is not a true loop as all elements will be execute at once
	ld s2, 24(sp)			# s2: W_I_1
	ld s3, 96(sp)			# s3: W_O_1
	ld s4, 48(sp)			# s4: W_F_1
	ld s5, 72(sp)			# s5: W_G_1

	mul s1, s11, a1			# m * hidenode_1
	add s1, s1, s10			# m * hidenode_1 + j
	slli s1, s1, 2			# 4 bytes
	
	add s2, s2, s1			# W_I_1[m][j]
	add s3, s3, s1			# W_O_1[m][j]
	add s4, s4, s1			# W_F_1[m][j]
	add s5, s5, s1			# W_G_1[m][j]
	mv s1, a1
	slli s1, s1, 2
	vsetvli t0, a0, e32, m2, ta, ma
	vle32.v v4, (t4)		# v4 = x[p][0~t0]
	vlse32.v v6, (s2), s1
	vlse32.v v8, (s3), s1
	vlse32.v v10, (s4), s1
	vlse32.v v12, (s5), s1

	vfmul.vv v6, v6, v4
	vfmul.vv v8, v8, v4
	vfmul.vv v10, v10, v4
	vfmul.vv v12, v12, v4

	vfredsum.vs v6, v6, v0
	vfredsum.vs v8, v8, v0
	vfredsum.vs v10, v10, v0
	vfredsum.vs v12, v12, v0

	vfmv.f.s f1, v6
	vfmv.f.s f2, v8
	vfmv.f.s f3, v10
	vfmv.f.s f4, v12

	add s11, s11, t0
	blt s11, a0, loop4		# m < innode(a0)




	li s11, 0                       # s11: m
        mv s1, a1
        slli s1, s1, 2                  # hidenode_1 * 4
        sub s6, t6, s1                  # s6: h_pre = H - hidenode_1 * 4
        sub s7, t5, s1                  # s7: state_pre = S - hidenode_1 * 4
loop5:
        ld s2, 32(sp)                   # s2: U_I_1
        ld s3, 104(sp)                  # s3: U_O_1
        ld s4, 56(sp)                   # s4: U_F_1
        ld s5, 80(sp)                   # s5: U_G_1

        mul s1, s11, a1                 # m * hidenode_1
        add s1, s1, s10                 # m * hidenode_1 + j
        slli s1, s1, 2                  # 4 bytes

        add s2, s2, s1                  # U_I_1[m][j]
        add s3, s3, s1                  # U_O_1[m][j]
        add s4, s4, s1                  # U_F_1[m][j]
        add s5, s5, s1                  # U_G_1[m][j]

        mv s1, a1
        slli s1, s1, 2                  # hidenode * 4

	vsetvli t0, a0, e32, m2, ta, ma
        vle32.v v4, (s6)
        vlse32.v v6, (s2), s1
        vlse32.v v8, (s3), s1
        vlse32.v v10, (s4), s1
        vlse32.v v12, (s5), s1

        vfmul.vv v6, v6, v4
        vfmul.vv v8, v8, v4
        vfmul.vv v10, v10, v4
        vfmul.vv v12, v12, v4

        vfredsum.vs v6, v6, v0
        vfredsum.vs v8, v8, v0
        vfredsum.vs v10, v10, v0
        vfredsum.vs v12, v12, v0

        vfmv.f.s f5, v6
        vfmv.f.s f6, v8
        vfmv.f.s f7, v10
        vfmv.f.s f8, v12
        fadd.s f1, f1, f5
        fadd.s f2, f2, f6
        fadd.s f3, f3, f7
        fadd.s f4, f4, f8
	
	add s11, s11, t0
        blt s11, a1, loop5              # m < hidenode_1(a1)


#------------------------------------- end of loop

	ld s2, 40(sp)
	ld s3, 112(sp)
	ld s4, 64(sp)
	ld s5, 88(sp)
	
	mv s1, s10
	slli s1, s1, 2			# offset = j * 4
	add s2, s2, s1
	add s3, s3, s1
	add s4, s4, s1
	add s5, s5, s1
	
	addi sp, sp, -16
	sd ra, 8(sp)	

	li s1, 1
	fcvt.s.w f9, s1			# f9 = 1.0
	
	flw f5, (s2)			# B_I_1[j]
	fadd.s f5, f5, f1		# B_I_1[j] + inGate
	fneg.s f5, f5			# f5 = -f5
	jal ra, calc			# f6 = exp(-x)
	fadd.s f6, f6, f9		# f6 += 1.0
	fdiv.s f1, f9, f6		# f1 = 1.0 / (1 + exp(-x))
	
	flw f5, (s3)			# B_O_1[j]
	fadd.s f5, f5, f2		# B_O_1[j] + outGate
	fneg.s f5, f5			# f5 = -f5
	jal ra, calc			# f6 = exp(-x)
	fadd.s f6, f6, f9		# f6 += 1.0
	fdiv.s f2, f9, f6		# f2 = 1.0 / (1 + exp(-x))
	

	flw f5, (s4)			# B_F_1[j]
	fadd.s f5, f5, f3		# B_F_1[j] + forgetGate
	fneg.s f5, f5			# f5 = -f5
	jal ra, calc			# f6 = exp(-x)
	fadd.s f6, f6, f9		# f6 += 1.0
	fdiv.s f3, f9, f6		# f3 = 1.0 / (1 + exp(-x))
	
	flw f5, (s5)			# B_G_1[j]
	fadd.s f5, f5, f4		# B_G_1[j] + gGate
	li s5, 2
	fcvt.s.w f7, s5			# f7 = 2.0
	fmul.s f5, f5, f7		# f5 = 2f5: 2x
	fle.s s5, f5, f0
	beqz s5, d_calc
	fneg.s f5, f5			# f5 = -f5
d_calc:					# calc directly
	jal ra, calc			# f6 = exp(2x)
	fsub.s f7, f6, f9		# f7 = exp(2x) - 1
	fadd.s f6, f6, f9		# f6 = exp(2x) + 1
	fdiv.s f4, f7, f6		# f4 = (exp(2x) - 1) / (exp(2x) + 1)
	beqz s5, go_on			# f4 = -(exp(-2x) - 1) / (exp(-2x) + 1)
	fneg.s f4, f4	

go_on:
	ld ra, 8(sp)
	addi sp, sp, 16	

#----------------------------- write

	mv s1, s10
	slli s1, s1, 2			# j * 4
	mv s2, a7
	ld s3, 0(sp)
	ld s4, 8(sp)
	ld s5, 16(sp)
	add s2, s2, s1			# in_gate_layer + j * 4
	add s3, s3, s1
	add s4, s4, s1
	add s5, s5, s1
	
	fsw f1, (s2)
	fsw f2, (s3)
	fsw f3, (s4)
	fsw f4, (s5)	
	


	mv s1, s10
	slli s1, s1, 2			# j * 4
	add s7, s7, s1			# &state_pre[j]
	flw f5, (s7)			# float s_pre = state_pre[j];
	fmul.s f5, f5, f3		# foget_gate_layer1[j] * s_pre
	fmul.s f6, f1, f4		# g_gate_layer1[j] * in_gate_layer1[j]
	fadd.s f5, f5, f6			
	add s7, t5, s1			# state[j]
	fsw f5, (s7)	

#------------------------------ calc h[j]
	addi sp, sp, -16
	sd ra, 8(sp)

	li s5, 2
        fcvt.s.w f7, s5                 # f7 = 2.0
        fmul.s f5, f5, f7               # f5 = 2f5: 2x
        fle.s s5, f5, f0
        beqz s5, d_calc_1
        fneg.s f5, f5                   # f5 = -f5
d_calc_1:                               # calc directly
        jal ra, calc                    # f6 = exp(2x)
        fsub.s f7, f6, f9               # f7 = exp(2x) - 1
        fadd.s f6, f6, f9               # f6 = exp(2x) + 1
        fdiv.s f4, f7, f6               # f4 = (exp(2x) - 1) / (exp(2x) + 1)
        beqz s5, go_on_1                # f4 = -(exp(-2x) - 1) / (exp(-2x) + 1)
        fneg.s f4, f4

go_on_1:

	ld ra, 8(sp)
	addi sp, sp, 16
#------------------------------
	
	fmul.s f5, f2, f4		# out_gate_layer1 * tanh(state[j])
	
	mv s1, s10
	slli s1, s1, 2
	add s7, t6, s1
	fsw f5, (s7)			# h[j] = f5


#------------------------------ loop
	addi s10, s10, 1
	blt s10, a1, loop3		# j < hidenode_1(a1)
	
	addi t1, t1, 1
	blt t1, a2, loop2		# p < time_step(a2)

#	fmv.s f5, f1
#	addi sp, sp, -16
#	sd ra, 8(sp)
#
#	jal ra, calc
#	
#	fsw f6, (s2)
#	
#	ld ra, 8(sp)
#	addi sp, sp, 16
	

# ------------------------------------- after 4 networks
	beqz a3, finished
	ld s5, 128(sp)			# s2: W_out
	ld s3, 136(sp)			# s3: B_out
					# t6: h
					# a3: predict

	vsetvli t0, a1, e32, m2, ta, ma	# vl = min(vl, hidenode_4)
	vle32.v v4, (t6)		# v4: h[0/1/2]

	li s7, 4
	mul s7, s7, a0			# 4 * (innode = outnode = 3)
	li s10, 0			# k: 0, outnode

last_loop:
	vlse32.v v6, (s5), s7		# v6: W[0/1/2, k]

	vfmul.vv v6, v6, v4		# h[0/1/2] * W[0/1/2, k]
	vfredsum.vs v6, v6, v0		# v[6,0] = mul_sum

	vfmv.f.s f1, v6			# out = mul_sum

#	fsw f1, (a3)
	flw f2, (s3)
	fadd.s f5, f1, f2

	li s4, 1
	fcvt.s.w f9, s4
	fneg.s f5, f5
# ----------------------------- 	calc exp
	addi sp, sp, -16
	sd ra, 8(sp)

	jal ra, calc
	fadd.s f6, f6, f9
	fdiv.s f7, f9, f6

	ld ra, 8(sp)
	addi sp, sp, 16
# ----------------------------
	fsw f7, (a3)

	addi s3, s3, 4
	addi s5, s5, 4
	addi a3, a3, 4
	addi s10, s10, 1
	blt s10, a1, last_loop	


finished:
	ret


# t3 s1 s6 s8 s9 s2(load后就没用了)
calc:					# f6 = exp(f5)
#	li s2, 0			# s2: signed bit
#	fcvt.s.w f8, s2
	fle.s s2, f5, f0
	beqz s2, unsigned
	fneg.s f5, f5

unsigned:
	li s1, 1
	slli s1, s1, 8
	fcvt.s.w f6, s1
	fmul.s f5, f5, f6
	fcvt.w.s s8, f5			# x << 8

	li t3, 0
	ld s6, 136(sp)			# s6: pow_e  Notice! sp has been changed since called calc
	li s1, 1
	fcvt.s.w f6, s1			# res = 1.0

loop_s:
	andi s9, s8, 1
	beqz s9, next
	li s1, 4
	mul s1, s1, t3
	add s1, s1, s6
	flw f7, (s1)
	fmul.s f6, f6, f7
next:
	srli s8, s8, 1
	addi t3, t3, 1
	li s1, 12
	blt t3, s1, loop_s

	beqz s2, return
	li s1, 1
	fcvt.s.w f8, s1
	fdiv.s f6, f8, f6		# res = 1.0 / res

return:
	jr ra 
